/**
 * 
 */
package edu.gatech.cc.liam.core;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map.Entry;
import java.util.Set;


/**
 * @author liam
 *
 */
public class ProbabilityMap<KeyType, EventType> {
	
	private static final long serialVersionUID = -6545204281005405504L;
	
	//theProps maps States and then a jointActionType to the distribution of next states
	final private HashMap<KeyType, DiscreteDistribution<EventType>> probMap;
	
	public ProbabilityMap() {
		this.probMap = new HashMap<KeyType, DiscreteDistribution<EventType>>();
	}	
	
	public void PutProb(KeyType key, EventType event, Double d) {
		DiscreteDistribution<EventType> theDist;
		if(probMap.containsKey(key)) {
			theDist = probMap.get(key);
		} else {
			theDist = new DiscreteDistribution<EventType>();
			probMap.put(key, theDist);
		}
		theDist.put(event, d);
	}

	public DiscreteDistribution<EventType> getDistribution(KeyType key) {
		assert (probMap.containsKey(key));
		return probMap.get(key);
	}
	public double getProb(KeyType key, EventType e) {
		return probMap.get(key).get(e);
	}

	public boolean ValidateAllProbs() {
		for (Entry<KeyType, DiscreteDistribution<EventType>> entry: probMap.entrySet()) {
			if(entry.getValue().isDistribution() == false) {
				System.out.println("**ERROR** The distribution " + entry.getValue() + " for key " + entry.getKey() + " is not a valid distribution!");
				return false;
			}
		}
		return true;
	}

	public void normalizeTransitionFunction() {
		for (DiscreteDistribution<EventType> e: probMap.values()) {
			e.normalize();
		}
	}
	
	@Override
	public String toString() {
		return probMap.toString();
	}
	public Set<EventType> getEventSet(KeyType key) {
		return this.probMap.get(key).keySet();
	}
	
}
